﻿using System;
using System.Buffers;
using System.Diagnostics;
using System.Runtime.CompilerServices;
using eLibrary;

namespace Letter.IO.Buf;

sealed class Buffer<T> : IBufferWriter<T>, IBuffer<T>
{
    private BufferSegmentStack<T> _bufferSegmentPool;
    private readonly BufferOptions<T> _options;

    private long _unconsumedBytes;
    private long _unflushedBytes;
    private long _lastExaminedIndex = -1;

    private BufferSegment<T> _readHead;
    private int _readHeadIndex;

    private bool _disposed;

    private BufferSegment<T> _readTail;
    private int _readTailIndex;
    private int _minimumReadBytes;

    private BufferSegment<T> _writingHead;
    private Memory<T> _writingHeadMemory;
    private int _writingHeadBytesBuffered;

    private int MinimumSegmentSize => _options.MinimumSegmentSize;
    private long PauseWriterThreshold => _options.PauseWriterThreshold;
    private long ResumeWriterThreshold => _options.ResumeWriterThreshold;

    private long Length => _unconsumedBytes;

    public Buffer(BufferOptions<T> options)
    {
        _bufferSegmentPool = new BufferSegmentStack<T>(options.InitialSegmentPoolSize);
        _options = options;
    }

    private void ResetState()
    {
        _readTailIndex = 0;
        _readHeadIndex = 0;
        _lastExaminedIndex = -1;
        _unflushedBytes = 0;
        _unconsumedBytes = 0;
    }

    private Memory<T> InternalGetMemory(int sizeHint)
    {
        if (_disposed)
        {
            throw new InvalidOperationException("The buffer has been disposed");
        }

        if (sizeHint < 0)
        {
            throw new ArgumentOutOfRangeException(ExceptionArgument.minimumSize.ToString());
        }

        AllocateWriteHeadIfNeeded(sizeHint);

        return _writingHeadMemory;
    }

    private Span<T> InternalGetSpan(int sizeHint)
    {
        return InternalGetMemory(sizeHint).Span;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private void AllocateWriteHeadIfNeeded(int sizeHint)
    {
        if (_writingHeadMemory.Length == 0 || _writingHeadMemory.Length < sizeHint)
        {
            AllocateWriteHeadSynchronized(sizeHint);
        }
    }

    private void AllocateWriteHeadSynchronized(int sizeHint)
    {
        if (_writingHead == null)
        {
            BufferSegment<T> newSegment = AllocateSegment(sizeHint);

            _writingHead = _readHead = _readTail = newSegment;
            _lastExaminedIndex = 0;
        }
        else
        {
            int bytesLeftInBuffer = _writingHeadMemory.Length;

            if (bytesLeftInBuffer == 0 || bytesLeftInBuffer < sizeHint)
            {
                if (_writingHeadBytesBuffered > 0)
                {
                    _writingHead.End += _writingHeadBytesBuffered;
                    _writingHeadBytesBuffered = 0;
                }

                if (_writingHead.Length == 0)
                {
                    _writingHead.ResetMemory();
                    RentMemory(_writingHead, sizeHint);
                }
                else
                {
                    BufferSegment<T> newSegment = AllocateSegment(sizeHint);

                    _writingHead.SetNext(newSegment);
                    _writingHead = newSegment;
                }
            }
        }
    }

    private BufferSegment<T> AllocateSegment(int sizeHint)
    {
        BufferSegment<T> newSegment = CreateSegmentUnsynchronized();

        RentMemory(newSegment, sizeHint);

        return newSegment;
    }

    private void RentMemory(BufferSegment<T> segment, int sizeHint)
    {
        // Segment should be new or reset, otherwise a memory leak could occur
        Debug.Assert(segment.MemoryOwner is null);
        Debug.Assert(sizeHint >= 0);

        MemoryPool<T> pool = null;
        int maxSize = -1;

        if (!_options.IsDefaultSharedMemoryPool)
        {
            pool = _options.Pool;
            maxSize = pool.MaxBufferSize;
        }

        if (sizeHint <= maxSize)
        {
            // Use the specified pool as it fits. Specified pool is not null as maxSize == -1 if _pool is null.
            segment.SetOwnedMemory(pool!.Rent(GetSegmentSize(sizeHint, maxSize)));
        }
        else
        {
            // Use the array pool
            int sizeToRequest = GetSegmentSize(sizeHint);
            segment.SetOwnedMemory(ArrayPool<T>.Shared.Rent(sizeToRequest));
        }

        _writingHeadMemory = segment.AvailableMemory;
    }

    private int GetSegmentSize(int sizeHint, int maxBufferSize = int.MaxValue)
    {
        // First we need to handle case where hint is smaller than minimum segment size
        sizeHint = Math.Max(MinimumSegmentSize, sizeHint);
        // After that adjust it to fit into pools max buffer size
        int adjustedToMaximumSize = Math.Min(maxBufferSize, sizeHint);
        return adjustedToMaximumSize;
    }

    private BufferSegment<T> CreateSegmentUnsynchronized()
    {
        if (_bufferSegmentPool.TryPop(out BufferSegment<T> segment))
        {
            return segment;
        }

        return new BufferSegment<T>();
    }

    private void ReturnSegmentUnsynchronized(BufferSegment<T> segment)
    {
        Debug.Assert(segment != _readHead, "Returning _readHead segment that's in use!");
        Debug.Assert(segment != _readTail, "Returning _readTail segment that's in use!");
        Debug.Assert(segment != _writingHead, "Returning _writingHead segment that's in use!");

        if (_bufferSegmentPool.Count < _options.MaxSegmentPoolSize)
        {
            _bufferSegmentPool.Push(segment);
        }
    }

    private bool Commit()
    {
        if (_unflushedBytes == 0)
        {
            // Nothing written to commit
            return false;
        }

        // Update the writing head
        Debug.Assert(_writingHead != null);
        _writingHead.End += _writingHeadBytesBuffered;

        // Always move the read tail to the write head
        _readTail = _writingHead;
        _readTailIndex = _writingHead.End;

        long oldLength = _unconsumedBytes;
        _unconsumedBytes += _unflushedBytes;

        bool resumeReader = true;

        if (_unconsumedBytes < _minimumReadBytes)
        {
            // Don't yield the reader if we haven't written enough
            resumeReader = false;
        }

        _unflushedBytes = 0;
        _writingHeadBytesBuffered = 0;

        return resumeReader;
    }

    private void InternalAdvance(int bytes)
    {
        if (_disposed)
        {
            throw new InvalidOperationException("The buffer has been disposed");
        }

        if ((uint)bytes > (uint)_writingHeadMemory.Length)
        {
            throw new ArgumentOutOfRangeException(ExceptionArgument.bytes.ToString());
        }

        AdvanceCore(bytes);
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private void AdvanceCore(int bytesWritten)
    {
        _unflushedBytes += bytesWritten;
        _writingHeadBytesBuffered += bytesWritten;
        _writingHeadMemory = _writingHeadMemory.Slice(bytesWritten);
    }

    private void InternalFlush()
    {
        if (_disposed)
        {
            throw new InvalidOperationException("The buffer has been disposed");
        }

        Commit();
    }

    private void InternalAdvanceReader(in SequencePosition consumed)
    {
        InternalAdvanceReader(consumed, consumed);
    }

    private void InternalAdvanceReader(in SequencePosition consumed, in SequencePosition examined)
    {
        if (_disposed)
        {
            throw new InvalidOperationException("The buffer has been disposed");
        }

        // TODO: Use new SequenceMarshal.TryGetReadOnlySequenceSegment to get the correct data
        // directly casting only works because the type value in ReadOnlySequenceSegment is 0
        InternalAdvanceReader((BufferSegment<T>)consumed.GetObject(), consumed.GetInteger(), (BufferSegment<T>)examined.GetObject(), examined.GetInteger());
    }

    private void InternalAdvanceReader(BufferSegment<T> consumedSegment, int consumedIndex, BufferSegment<T> examinedSegment, int examinedIndex)
    {
        if (consumedSegment != null && examinedSegment != null && BufferSegment<T>.GetLength(consumedSegment, consumedIndex, examinedSegment, examinedIndex) < 0)
        {
            throw new InvalidOperationException("The examined position must be greater than or equal to the consumed position.");
        }

        BufferSegment<T> returnStart = null;
        BufferSegment<T> returnEnd = null;

        var examinedEverything = false;
        if (examinedSegment == _readTail)
        {
            examinedEverything = examinedIndex == _readTailIndex;
        }

        if (examinedSegment != null && _lastExaminedIndex >= 0)
        {
            long examinedBytes = BufferSegment<T>.GetLength(_lastExaminedIndex, examinedSegment, examinedIndex);
            long oldLength = _unconsumedBytes;

            if (examinedBytes < 0)
            {
                throw new InvalidOperationException("The examined position cannot be less than the previously examined position.");
            }

            _unconsumedBytes -= examinedBytes;

            // Store the absolute position
            _lastExaminedIndex = examinedSegment.RunningIndex + examinedIndex;

            Debug.Assert(_unconsumedBytes >= 0, "Length has gone negative");
            Debug.Assert(ResumeWriterThreshold >= 1, "ResumeWriterThreshold is less than 1");

        }

        if (consumedSegment != null)
        {
            if (_readHead == null)
            {
                throw new InvalidOperationException("The PipeReader has already advanced past the provided position.");
            }

            returnStart = _readHead;
            returnEnd = consumedSegment;

            void MoveReturnEndToNextBlock()
            {
                BufferSegment<T> nextBlock = returnEnd!.NextSegment;
                if (_readTail == returnEnd)
                {
                    _readTail = nextBlock;
                    _readTailIndex = 0;
                }

                _readHead = nextBlock;
                _readHeadIndex = 0;

                returnEnd = nextBlock;
            }

            if (consumedIndex == returnEnd.Length)
            {
                // If the writing head isn't block we're about to return, then we can move to the next one
                // and return this block safely
                if (_writingHead != returnEnd)
                {
                    MoveReturnEndToNextBlock();
                }
                // If the writing head is the same as the block to be returned, then we need to make sure
                // there's no pending write and that there's no buffered data for the writing head
                else if (_writingHeadBytesBuffered == 0)//if (_writingHeadBytesBuffered == 0 && !_operationState.IsWritingActive)
                {
                    // Reset the writing head to null if it's the return block and we've consumed everything
                    _writingHead = null;
                    _writingHeadMemory = default;

                    MoveReturnEndToNextBlock();
                }
                else
                {
                    _readHead = consumedSegment;
                    _readHeadIndex = consumedIndex;
                }
            }
            else
            {
                _readHead = consumedSegment;
                _readHeadIndex = consumedIndex;
            }
        }

        while (returnStart != null && returnStart != returnEnd)
        {
            BufferSegment<T> next = returnStart.NextSegment;
            returnStart.Reset();
            ReturnSegmentUnsynchronized(returnStart);
            returnStart = next;
        }
    }
    
    private ReadOnlySequence<T> InternalRead()
    {
        if (_disposed)
        {
            throw new InvalidOperationException("The buffer has been disposed");
        }

        ReadOnlySequence<T> result;
        // If the awaitable is already complete then return the value result directly
        GetReadResult(out result);

        return result;
    }
    
    private void InternalRelease()
    {
        if (_disposed)
        {
            return;
        }

        _disposed = true;
        // Return all segments
        // if _readHead is null we need to try return _commitHead
        // because there might be a block allocated for writing
        BufferSegment<T> segment = _readHead ?? _readTail;
        while (segment != null)
        {
            BufferSegment<T> returnSegment = segment;
            segment = segment.NextSegment;

            returnSegment.Reset();
        }

        _writingHead = null;
        _writingHeadMemory = default;
        _readHead = null;
        _readTail = null;
        _lastExaminedIndex = -1;
    }

    private void GetReadResult(out ReadOnlySequence<T> result)
    {
        // No need to read end if there is no head
        BufferSegment<T> head = _readHead;
        if (head != null)
        {
            Debug.Assert(_readTail != null);
            // Reading commit head shared with writer
            var readOnlySequence = new ReadOnlySequence<T>(head, _readHeadIndex, _readTail, _readTailIndex);
            result = readOnlySequence;
        }
        else
        {
            result = default;
        }

        // Reset the minimum read bytes when read yields
        _minimumReadBytes = 0;
    }


    private void InternalReset()
    {
        if (!_disposed)
        {
            throw new InvalidOperationException("Both reader and writer has to be completed to be able to reset the pipe.");
        }

        _disposed = false;

        ResetState();
    }

    
    object IBuffer<T>.Token
    {
        get;
        set;
    }

    Memory<T> IBufferWriter<T>.GetMemory(int sizeHint)
    {
        return InternalGetMemory(sizeHint);
    }

    Span<T> IBufferWriter<T>.GetSpan(int sizeHint)
    {
        return InternalGetSpan(sizeHint);
    }
    
    void IBufferWriter<T>.Advance(int count)
    {
        InternalAdvance(count);
    }
    
    long IReadableBuffer<T>.ReadableLength
    {
        get { return this.Length; }
    }

    ReadOnlySequence<T> IReadableBuffer<T>.GetReadableSequence()
    {
        return this.InternalRead();
    }

    void IReadableBuffer<T>.ReaderAdvanceTo(SequencePosition consumed)
    {
        this.InternalAdvanceReader(consumed);
    }

    void IReadableBuffer<T>.ReaderAdvanceTo(SequencePosition consumed, SequencePosition examined)
    {
        this.InternalAdvanceReader(consumed, examined);
    }

    Memory<T> IWritableBuffer<T>.GetWritableSequence(int size)
    {
        return this.InternalGetMemory(size);
    }

    void IWritableBuffer<T>.WriterAdvance(int count)
    {
        this.InternalAdvance(count);
    }

    void IWritableBuffer<T>.Flush()
    { 
        this.InternalFlush();
    }

    void IClearable.Clear()
    {
        var buf = this.InternalRead();
        this.InternalAdvanceReader(buf.End);
    }

    void IReleasable.Release()
    {
        var clearable = (IClearable) this;
        clearable.Clear();
        this.InternalRelease();
    }

    void IResettable.Reset()
    {
        this.InternalReset();
    }
}